package cn.monkey.state.core.reactive;

import cn.monkey.state.core.FetchStateGroup;
import cn.monkey.state.core.StateGroup;
import cn.monkey.state.bean.Refreshable;
import reactor.core.publisher.Mono;

import java.lang.invoke.MethodHandles;
import java.lang.invoke.VarHandle;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class SimpleReactiveStateGroupPool implements ReactiveStateGroupPool, Refreshable {
    protected final ReactiveStateGroupFactory stateGroupFactory;
    protected volatile Map<String, StateGroup> stateGroupMap;
    static final VarHandle VAR_HANDLE;

    static {
        MethodHandles.Lookup lookup = MethodHandles.lookup();
        try {
            VAR_HANDLE = lookup.findVarHandle(SimpleReactiveStateGroupPool.class, "stateGroupMap", Map.class);
        } catch (NoSuchFieldException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    public SimpleReactiveStateGroupPool(ReactiveStateGroupFactory stateGroupFactory) {
        this.stateGroupFactory = stateGroupFactory;
        this.stateGroupMap = new ConcurrentHashMap<>();
    }

    @Override
    public Mono<FetchStateGroup> findOrCreate(String id, Object... args) {
        return Mono.justOrEmpty(this.stateGroupMap.get(id))
                .map(stateGroup -> new FetchStateGroup(false, stateGroup))
                .switchIfEmpty(Mono.defer(() -> this.create(id, args)));
    }

    Mono<FetchStateGroup> create(String id, Object... args) {
        return Mono.justOrEmpty(this.stateGroupMap.get(id))
                .flatMap(stateGroup -> {
                    if (stateGroup.canClose()) {
                        return this.stateGroupFactory.create(id, args)
                                .map(sg -> new FetchStateGroup(true, this.stateGroupMap.computeIfPresent(id, (k, v) -> sg)));
                    }
                    return Mono.just(new FetchStateGroup(false, stateGroup));
                })
                .switchIfEmpty(Mono.defer(() -> this.stateGroupFactory.create(id, args)
                        .map(stateGroup -> new FetchStateGroup(true, stateGroup))));
    }

    @Override
    public void refresh() {
        final Map<String, StateGroup> stateGroupMap = this.stateGroupMap;
        ConcurrentHashMap<String, StateGroup> newStateGroupMap = new ConcurrentHashMap<>();
        for (Map.Entry<String, StateGroup> e : stateGroupMap.entrySet()) {
            StateGroup value = e.getValue();
            if (value.canClose()) {
                value.close();
                continue;
            }
            newStateGroupMap.put(e.getKey(), e.getValue());
        }
        VAR_HANDLE.compareAndSet(this, stateGroupMap, newStateGroupMap);
    }
}
